import os
from openai import AzureOpenAI
#Note: The openai-python library support for Azure OpenAI is in preview.
#Note: This code sample requires OpenAI Python library version 1.0.0 or higher.
import json
import time
import shutil
import requests
import datetime
import multiprocessing
import pandas as pd
import multiprocessing as mp
from docx import Document as document
from docx.document import Document

import tiktoken

chatgpt_encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
gpt_4_encoding = tiktoken.encoding_for_model("gpt-4")


def json_to_text(json_path):
    '''读取json文件并返回字符串

    '''
    with open(json_path, 'r') as file:
        json_data = json.load(file)

    total_text = '\n'.join(list(json_data['knowledges'].values()))

    if "abstract" in json_data:
        if len(json_data["abstract"]) > 20:
            # 存在摘要时返回False
            return False, total_text, json_data["abstract"]

    return True, total_text, ''  # 返回（是否需要生成摘要状态，全文，摘要）

def write_abstract_use_chatgpt_by_azure(content_dict_list,
                                        files,
                                        json_dir='./abstract_for_json_file',
                                        model_deployment_name_on_azure="GPT4",
                                        do_check=False):
    '''

    https://portal.azure.com/
    # 账号密码请联系陈艺荣

    model_deployment_name_on_azure="GPT35" or "GPT4"
    # 需要确定具体的部署版本号请联系陈艺荣
    # GPT4: gpt-4-1106-Preview, 输入：128,000 输出：4,096, 训练数据（上限）：2023年4月
    # GPT35: gpt-35-turbo-1106, 输入：16,385 输出：4,096, 训练数据（上限）：2021年9月
    '''
    if model_deployment_name_on_azure == "GPT4":
        # now = time.localtime()
        # current_date = time.strftime("%Y-%m", now)
        # system_prompt_zh = f'你是一个智能助手，一个经过大规模语料训练与价值对齐的语言模型。\n知识截止: 2021-09\n当前日期: {current_date}'
        system_prompt = f'You are ChatGPT, a large language model trained by OpenAI.\nKnowledge cutoff: 2023-04\nCurrent date: 2024-03'
        client = AzureOpenAI(
            azure_endpoint="https://zhishenggpt40.openai.azure.com/",
            api_key='126b11dac9994145a291b9f0a37e53d7',  # 机密文件，在.bashrc中设置了系统变量
            api_version="2024-02-15-preview"
        )
    for key in range(len(content_dict_list)):
        temp_dict = content_dict_list[key]
        if temp_dict['file_name'] not in files:
            # old_time放在程序运行开始的地方
            old_time = time.time()

            content = temp_dict['content']
            if len(chatgpt_encoding.encode(content)) > 50000:
                content = content[0:50000]

            message_text = [{"role": "system", "content": system_prompt},
                            {"role": "user", "content": '''给定以下文档:\n\n''' + content + '''\n\n请依据上述文档内容，撰写不超过500字的摘要：'''}
                            ]

            completion = client.chat.completions.create(
                model=model_deployment_name_on_azure,  # model = "deployment_name"
                messages=message_text,
                temperature=0.7,
                max_tokens=4096,
                top_p=0.95,
                frequency_penalty=0,
                presence_penalty=0,
                stop=None
            )

            abstract = completion.choices[0].message.content

            temp_dict['chatgpt_abstract'] = abstract

            with open(os.path.join(json_dir, temp_dict['file_name'] + '.json'), "w", encoding="utf-8") as f:
                json.dump(temp_dict, f, indent=4, ensure_ascii=False)

            print(temp_dict['file_name'] + '.json' + " 保存成功！")

            current_time = time.time()

            if (current_time - old_time) < 20:
                time.sleep(35)
            elif (current_time - old_time) < 30:
                time.sleep(25)
            elif (current_time - old_time) < 40:
                time.sleep(15)
            elif (current_time - old_time) < 50:
                time.sleep(5)


class Demo(multiprocessing.Process):
    def __init__(self, idx_process, dict_list, files, num_process=32, save_dir='./abstract_for_json_file_20240304',
                 model_deployment_name_on_azure="GPT4", do_check=False):
        '''
        idx_process: 线程号
        knowledge_dict_list: 待转换的知识列表
        files: 已经转换的文件列表

        '''
        multiprocessing.Process.__init__(self)
        self.idx_process = idx_process
        self.dict_list = dict_list
        self.files = files
        self.num_process = num_process
        self.save_dir = save_dir
        self.model_deployment_name_on_azure = model_deployment_name_on_azure
        self.do_check = do_check

    def run(self):
        # 启动进程
        print(f'【进程{self.idx_process}】' + str(datetime.datetime.now()))

        write_abstract_use_chatgpt_by_azure(content_dict_list=self.dict_list,
                                            files=self.files,
                                            json_dir=self.save_dir,
                                            model_deployment_name_on_azure=self.model_deployment_name_on_azure,
                                            do_check=False)


def main():
    # 运行进程数
    num_process = 10
    model_deployment_name_on_azure = "GPT4"
    # GPT4: gpt-4-1106-preview 输入：128,000 输出：4,096
    # GPT35: gpt-3.5-turbo-1106 输入：16,385 输出：4,096
    # https://learn.microsoft.com/zh-cn/azure/ai-services/openai/concepts/models

    # txt_dir = '/home/mas-lin.huadong/projects/01-基因云_检索问答/docs/txt'
    input_json_dir = 'docs11111/json'
    json_dir = 'docs11111/abstract'

    if not os.path.exists(json_dir):
        os.makedirs(json_dir)

    # 读取已经转换的文件名
    files = os.listdir(json_dir)
    files = [os.path.splitext(file)[0] for file in files]

    # txt_files = os.listdir(txt_dir)
    input_json_files = os.listdir(input_json_dir)
    total_dict_list = []
    for file in input_json_files:
        state, content, abstract = json_to_text(json_path=os.path.join(input_json_dir, file))
        if state == False:
            # 将文件复制到json_dir

            temp_dict = {
                'file_name': os.path.splitext(file)[0],
                'json_name': file,
                'content': content,
                "chatgpt_abstract": abstract
            }
            if not os.path.exists(os.path.join(json_dir, temp_dict['file_name'] + '.json')):
                # 不存在文件则保存
                with open(os.path.join(json_dir, temp_dict['file_name'] + '.json'), "w", encoding="utf-8") as f:
                    json.dump(temp_dict, f, indent=4, ensure_ascii=False)

                print(temp_dict['file_name'] + '.json' + " 保存成功！")

        else:
            total_dict_list.append(
                {
                    'file_name': os.path.splitext(file)[0],
                    'json_name': file,
                    'content': content
                }
            )

    # 过滤掉已经转换的
    rawlist = total_dict_list  # 任务列表
    old_rawlist = total_dict_list  # 任务列表
    # 忽略已处理列表
    rawlist = [example for example in old_rawlist if str(example['file_name']) not in files]

    print("待处理数据个数：", len(rawlist))

    # 根据进程数切割任务
    temp = []
    id_list = []
    step = len(rawlist) // num_process
    for i in range(0, num_process - 1):
        temp.append(rawlist[step * i:step * i + step])
        id_list.append([str(example['file_name']) for example in rawlist[step * i:step * i + step]])
    temp.append(rawlist[step * (num_process - 1):])
    id_list.append([str(example['file_name']) for example in rawlist[step * (num_process - 1):]])
    # print(temp)

    # 初始化进程

    print("【启动】初始化进程...")

    processes = []
    for i in range(0, num_process):
        try:
            # 重新读取已经转换的文件名
            files = os.listdir(json_dir)
            files = [os.path.splitext(file)[0] for file in files]
            process = Demo(i, temp[i], files, num_process, save_dir=json_dir,
                           model_deployment_name_on_azure=model_deployment_name_on_azure)
            process.start()
            processes.append(process)
        except Exception as e:
            print('启动【进程{i}】出错，5秒后重试...')
            time.sleep(5)
            continue

    while len(files) < len(old_rawlist):
        for i, p in enumerate(processes):
            if not p.is_alive():
                processes.remove(p)
                print('-------------------------------------------')
                print(f"【进程{i}】正在重启...")
                time.sleep(30)
                # 重新读取已经转换的文件名
                files = os.listdir(json_dir)
                files = [os.path.splitext(file)[0] for file in files]
                # 重新启动进程
                if (set(id_list[i]) < set(files)):
                    # 列表元素已经全部转换，不重启进程
                    print(f"【进程{i}】已经完成任务，不再重启...")
                else:
                    process = Demo(i, temp[i], files, num_process, save_dir=json_dir,
                                   model_deployment_name_on_azure=model_deployment_name_on_azure)
                    process.start()
                    processes.append(process)

if __name__ == "__main__":

    try:
        start = time.time()
        main()
        end = time.time()
        print("运行程序花费了%s秒" % (end - start))
    except Exception as e:
        print(e)